package main

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

var ans []int

func distanceK(root *TreeNode, target *TreeNode, k int) []int {
	ans = nil
	parents := buildParent(root)
	goDown(target, nil, k)
	d := k - 1
	n := target
	for d >= 0 && n != nil {
		goDown(parents[n], n, d)
		n = parents[n]
		d--
	}
	return ans
}

// 向下搜索
func goDown(node *TreeNode, exclude *TreeNode, k int) {
	if node == nil || node == exclude {
		return
	}
	if k == 0 {
		ans = append(ans, node.Val)
		return
	}
	if node.Left != nil {
		goDown(node.Left, exclude, k-1)
	}
	if node.Right != nil {
		goDown(node.Right, exclude, k-1)
	}
}

func buildParent(root *TreeNode) map[*TreeNode]*TreeNode {
	m := make(map[*TreeNode]*TreeNode)
	var que []*TreeNode
	var n *TreeNode
	que = append(que, root)
	for len(que) > 0 {
		n = que[0]
		que = que[1:]
		if n.Left != nil {
			m[n.Left] = n
			que = append(que, n.Left)
		}
		if n.Right != nil {
			m[n.Right] = n
			que = append(que, n.Right)
		}
	}
	return m
}

func main() {
	tree := &TreeNode{
		Val: 1,
		Left: &TreeNode{
			Val: 2,
			Left: &TreeNode{
				Val: 4,
				Left: &TreeNode{
					Val: 8,
					Left: &TreeNode{
						Val: 10,
					},
				},
				Right: &TreeNode{
					Val: 9,
				},
			},
			Right: &TreeNode{
				Val: 5,
			},
		},
		Right: &TreeNode{
			Val: 3,
			Left: &TreeNode{
				Val: 6,
			},
			Right: &TreeNode{
				Val: 7,
			},
		},
	}
	m := buildParent(tree)
	println(m)
	goDown(tree, nil, 2)
	println(ans)
	distanceK(tree, tree.Left.Left, 2)
	println(ans)

}
